#include <filesystem>
#include "ColMap.h"
#include "PointIO.h"
#include "TensorMath.h"

namespace fs = std::filesystem;
using namespace torch::indexing;

namespace cm {

	//******************************************************************
	//FUNCTION:
	SInputData inputDataFromColmap(const std::string& vProjectRoot) {
		SInputData Ret;
		fs::path CmRoot(vProjectRoot);

		if (!fs::exists(CmRoot / "cameras.bin") && fs::exists(CmRoot / "sparse" / "0" / "cameras.bin")) {
			CmRoot = CmRoot / "sparse" / "0";
		}

		fs::path CamerasPath = CmRoot / "cameras.bin";
		fs::path ImagesPath = CmRoot / "images.bin";
		fs::path PointsPath = CmRoot / "points3D.bin";

		if (!fs::exists(CamerasPath)) throw std::runtime_error(CamerasPath.string() + " does not exist");
		if (!fs::exists(ImagesPath)) throw std::runtime_error(ImagesPath.string() + " does not exist");
		if (!fs::exists(PointsPath)) throw std::runtime_error(PointsPath.string() + " does not exist");

		std::ifstream Camf(CamerasPath.string(), std::ios::binary);
		if (!Camf.is_open()) throw std::runtime_error("Cannot open " + CamerasPath.string());
		std::ifstream Imgf(ImagesPath.string(), std::ios::binary);
		if (!Imgf.is_open()) throw std::runtime_error("Cannot open " + ImagesPath.string());

		size_t NumCameras = readBinary<uint64_t>(Camf);
		std::vector<CCamera> Cameras(NumCameras);

		std::unordered_map<uint32_t, CCamera*> CamMap;

		for (size_t i = 0; i < NumCameras; i++) {
			CCamera* pCam = &Cameras[i];
			pCam->setID(readBinary<uint32_t>(Camf));

			CameraModel Model = static_cast<CameraModel>(readBinary<int>(Camf)); // model ID
			pCam->setWidth(readBinary<uint64_t>(Camf));
			pCam->setHeight(readBinary<uint64_t>(Camf));

			if (Model == ESimplePinhole) {
				pCam->setFx(readBinary<double>(Camf));
				pCam->setFy(pCam->getFx());
				pCam->setCx(readBinary<double>(Camf));
				pCam->setCy(readBinary<double>(Camf));
			}
			else if (Model == EPinhole) {
				pCam->setFx(readBinary<double>(Camf));
				pCam->setFy(readBinary<double>(Camf));
				pCam->setCx(readBinary<double>(Camf));
				pCam->setCy(readBinary<double>(Camf));
			}
			else if (Model == ESimpleRadial) {
				pCam->setFx(readBinary<double>(Camf));
				pCam->setFy(pCam->getFx());
				pCam->setCx(readBinary<double>(Camf));
				pCam->setCy(readBinary<double>(Camf));
				pCam->setK1(readBinary<double>(Camf));
			}
			else if (Model == EOpenCV) {
				pCam->setFx(readBinary<double>(Camf));
				pCam->setFy(readBinary<double>(Camf));
				pCam->setCx(readBinary<double>(Camf));
				pCam->setCy(readBinary<double>(Camf));
				pCam->setK1(readBinary<double>(Camf));
				pCam->setK2(readBinary<double>(Camf));
				pCam->setP1(readBinary<double>(Camf));
				pCam->setP2(readBinary<double>(Camf));
			}
			else {
				throw std::runtime_error("Unsupported camera model: " + std::to_string(Model));
			}

			CamMap[pCam->getID()] = pCam;
		}

		Camf.close();


		size_t NumImages = readBinary<uint64_t>(Imgf);
		torch::Tensor UnorientedPoses = torch::zeros({ static_cast<long int>(NumImages), 4, 4 }, torch::kFloat32);

		for (size_t i = 0; i < NumImages; i++) {
			readBinary<uint32_t>(Imgf); // imageId

			torch::Tensor QVec = torch::tensor({
				readBinary<double>(Imgf),
				readBinary<double>(Imgf),
				readBinary<double>(Imgf),
				readBinary<double>(Imgf)
				}, torch::kFloat32);
			torch::Tensor R = quatToRotMat(QVec);
			torch::Tensor T = torch::tensor({
				{ readBinary<double>(Imgf) },
				{ readBinary<double>(Imgf) },
				{ readBinary<double>(Imgf) }
				}, torch::kFloat32);

			torch::Tensor Rinv = R.transpose(0, 1);
			torch::Tensor Tinv = torch::matmul(-Rinv, T);

			uint32_t CamId = readBinary<uint32_t>(Imgf);

			CCamera pCam = *CamMap[CamId];

			char Ch = '\0';
			std::string FilePath = "";
			while (true) {
				Imgf.read(&Ch, 1);
				if (Ch == '\0') break;
				FilePath += Ch;
			}

			// TODO: should "images" be an option?
			pCam.setFilePath((fs::path(vProjectRoot) / "images" / FilePath).string());

			UnorientedPoses[i].index_put_({ Slice(None, 3), Slice(None, 3) }, Rinv);
			UnorientedPoses[i].index_put_({ Slice(None, 3), Slice(3, 4) }, Tinv);
			UnorientedPoses[i][3][3] = 1.0f;

			// Convert COLMAP's camera CRS (OpenCV) to OpenGL
			UnorientedPoses[i].index_put_({ Slice(0, 3), Slice(1,3) }, UnorientedPoses[i].index({ Slice(0, 3), Slice(1,3) }) * -1.0f);

			size_t NumPoints2D = readBinary<uint64_t>(Imgf);
			for (size_t j = 0; j < NumPoints2D; j++) {
				readBinary<double>(Imgf); // x
				readBinary<double>(Imgf); // y
				readBinary<uint64_t>(Imgf); // point3D ID
			}

			Ret._Cameras.push_back(pCam);
		}

		Imgf.close();

		auto r = autoScaleAndCenterPoses(UnorientedPoses);
		torch::Tensor Poses = std::get<0>(r);
		Ret._Translation = std::get<1>(r);
		Ret._Scale = std::get<2>(r);
		Ret._Extend = (1 / Ret._Scale) * 1.1f;

		for (size_t i = 0; i < Ret._Cameras.size(); i++) {
			Ret._Cameras[i].setCamToWorld(Poses[i]);
		}

		CPointSet* pSet = readPointSet(PointsPath.string());
		torch::Tensor points = pSet->pointsTensor().clone();

		Ret._Points._XYZ = (points - Ret._Translation) * Ret._Scale;
		Ret._Points._RGB = pSet->colorsTensor().clone();

		RELEASE_POINTSET(pSet);

		return Ret;
	}

}